from utils.base_parser import BasicParser
import numpy as np


class JobDAG(object):

    def __init__(self, file_name, is_matrix=False):
        super(JobDAG, self).__init__()
        self._parse_from_xml(file_name, is_matrix)
        # 从xml解析得到dag结构

    def _parse_from_xml(self, file_name, is_matrix):
        parser = BasicParser(file_name)

        task_list = parser.generate_task_list()

        self.task_number = len(task_list)
        self.dependency = np.zeros((self.task_number, self.task_number))
        self.task_list = []

        self.pre_task_sets = []
        self.task_finish_time = [0] * self.task_number
        self.edge_set = []

        for _ in range(self.task_number):
            self.pre_task_sets.append(set([]))
        # add task list to
        self.add_task_list(task_list)

        dependencies = parser.generate_dependency()

        for pair in dependencies:
            self.add_dependency(pair[0], pair[1], pair[2])

    def add_task_list(self, task_list):
        self.task_list = task_list

        for i in range(0, len(self.task_list)):
            self.dependency[i][i] = task_list[i].running_time

    def add_dependency(self, pre_task_index, succ_task_index, transmission_cost):
        self.dependency[pre_task_index][succ_task_index] = transmission_cost
        self.pre_task_sets[succ_task_index].add(pre_task_index)

        # for each edge, we use a five dimension vector to represent this
        edge = [pre_task_index,
                self.task_list[pre_task_index].depth,
                self.task_list[pre_task_index].running_time,
                transmission_cost,
                succ_task_index,
                self.task_list[succ_task_index].depth,
                self.task_list[succ_task_index].running_time]

        self.edge_set.append(edge)